from ssummolib import load_index
from dict_to_phyloxml import write_xml
from count_hmms import countseqs
import os, re, sys, time, pickle
global results_file
global ambiguous_count
global hmmerdir
global blastdir
global blastdbdir
global top
ambiguous_count = 0
global results          ### This shall be the taxonomic dictionary with all the results in.
			### It shall represent tdict, but only includes nodes that have best hits.
			### For each node, keep accessions and scores.
			### e.g. results[node] = {'accessions' : [acc1, ... ] , 'subnodes' : { ... } }


###############################################################

################ User configured Variables ####################

hmmerdir = '/Users/albl500/bin'
blastdir = '/Users/albl500/bin'
blastdbdir = '/Users/albl500/454RDPstuff/blastDB'
blastversion = 2.6
top = '/Users/albl500/454RDPstuff/arbDB'

###############################################################

def parse_hmm(score_dict):
	"""This will parse the tabular result file from hmmsearch (must 
	specify --tblout txt ) and return the Evalue and
	the score.
	"""
#	x = re.compile('\s[0-9\.]+[e\-]*[0-9]*')
	x = re.compile('[a-zA-Z0-9\.\-]+\s+')
	accessions = set()
	with file('temp_table.txt','rU') as handle:
		for line in handle.xreadlines():
			score_list = []
			if line.startswith('#'):
				continue
			else:
				for result in x.finditer(line):
					result = result.group().strip()
					score_list.append(result)
				accession = score_list[0].strip()
				if '|' in accession:
					accession = accession[accession.find('|') + 1:]
				if accession in accessions:
					print "{0} appears twice in here".format(accession)
					print "Evalue results: 1st {0} 2nd {1}".format(score_dict[accession][0][-1],score_list[4])
					continue
				accessions.add(accession)
				if accession not in score_dict.keys():
					score_dict.update( { accession : [[score_list[4]] ,[ score_list[5] ]] } )
				else:
					score_dict[accession][0].append( score_list[4] ) # Evalues
					score_dict[accession][1].append( score_list[5] ) # Scores
	#os.system('rm -f temp_table.txt')
	return score_dict

def parse_pressed_hmm(results_handle, nodes):
	"""Given score_dict and a file handle (or name) for the results
	from hmmscan, this will update score_dict to contain the new scores.
	"""
	reg = re.compile('\s+')
	if type(results_handle) == str:
		handle = file(results_handle,'rU')
	elif type(results_handle) == file:
		handle = results_handle
	else:
		raise IOError("Wrong type of results_handle")
	temp_scores = {}  # { accession : [[Evalues, ... ] , [Scores, ... ]] , ... }
	for line in handle:
		if line.startswith('#'):
			continue
		results = reg.split(line)  # Splits the line by whitespace. Iterate through line entries.
		target = results[0]
		accession = results[2]
		Evalue = float(results[4])
		score =  float(results[5])
		found = 0
		if accession in temp_scores.keys():
			found = 1
			temp_scores['order'].append(target)
			temp_scores[accession][0].append(Evalue)
			temp_scores[accession][1].append(score)
		else:
			if found == 1:
				assert temp_scores.pop('order') == nodes			
			temp_scores.update({ 'order' : [target], accession : [ [Evalue], [score] ] } )
	del(temp_scores['order'])
	return temp_scores


def diverge( path, accession, choice_scores, tax_node ):
	"""Given a path, an accession and multiple best scores, go into each of them
	finding the overall best match"""
	global ambiguous_count
	global hmmerdir
	global blastdbdir
	global blastdir
	global blastversion
	global fastacmd
	if blastversion < 2.6:
		os.system('{0}/fastacmd -d {1}/tempDB -s {2} -o temp_seq.fa'.format(blastdir,blastdbdir,accession))
	elif blastversion == 2.6:
		os.system('{0}/blastdbcmd -db {1}/tempDB -entry {2} -out temp_seq.fa'.format(blastdir,blastdbdir,accession))
	else:
		print "Haven't tested on a blast version this recent. Attempting to use v2.6 config settings"
		os.system('{0}/blastdbcmd -db {1}/tempDB -entry {2} -out temp_seq.fa'.format(blastdir,blastdbdir,accession))
	score_dict = choice_scores.deepcopy()
	for OTU in sorted(score_dict.keys()):
		temp_path = path +'/' + OTU
		#choice_scores.update( {OTU :  {} })
		node = tax_node
		for dir in OTU.split('/'):
			node = node[dir]
		nodes = sorted(node.keys())
		### if nodes != []
		for next_OTU in nodes:
			if OTU + '/' + next_OTU.strip() in score_dict.keys():
				continue	## May have already been done.
			else:
				os.system( '{0}/hmmsearch -o /dev/null --tblout temp_table.txt --noali "{1}/{2}/{3}.hmm" temp_seq.fa'.format(hmmerdir, temp_path, next_OTU, next_OTU) )
				results = parse_hmm(score_dict)
				score_dict.update( {OTU + '/' + next_OTU.strip() : results[accession]})
	if score_dict == choice_scores:
		ambiguous_count += 1
		#Return ambiguous result #
		del(score_dict)
		return None
	else:
		best_score = 0
	for OTU in sorted(score_dict.keys()):
		OTU_score = max(score_dict[OTU][1])
		if OTU_score > best_score:
			best_score = OTU_score
			best_OTU = OTU
		else:
			continue
	best_count = 0
	best_dict = {}
	for item in score_dict.items():
		if best_score in item[1][1]:
			best_count += 1
			best_dict.update( { item[0] : best_score } )
		else:
			continue
	if best_count > 1:
		return diverge(path, accession, best_dict, tax_node)
	else:
		for OTU_score in choice_scores.items():
			if best_score == OTU_score[1]:
				best_OTU = OTU_score[0]
				return {accession : best_OTU[:best_OTU.find('/')]}
	

def score(score_dict, tax_node, path):
	global ambiguous_count
	global results_file
	best_key = None
	tax_node_keys = sorted(tax_node.keys())
	accessions_best = {}
	Evalue = 10
	score = 1
	#accessions = sorted (score_dict.keys())
	num_nodes = len(tax_node_keys)
	##### FIND BEST MATCHES, ASSIGNING ACCESSIONS TO NODES #####
	i = 0
	for accession in sorted(score_dict.keys()):
		i += 1
		diverge_set = set()
		Evalues = score_dict[accession][0]
		scores = score_dict[accession][1]
		key_len = len(Evalues)
		try:
			assert key_len == len(scores) and key_len == len(tax_node_keys)
		except AssertionError:
			print "{0}th iteration".format(i)
			print "Problem with accession {0}".format(accession)
			print "Numbers (lengths):-\n\tkey_len: {0}\tkey_scores: {1}\tnum_nodes: {2}".format(key_len,len(scores),num_nodes)
			print "Data:\n\tkey_scores: {0}\tNodes: {1}".format(scores,tax_node_keys)
			print "Score_dict keys: {0}   --- Items: {1} ".format( ' '.join([x for x in sorted(score_dict.keys())] ) ,' '.join( [ x for x in score_dict.values() ]) ) 
			raise
		
		#####################
		Evalues = score_dict[accession][0]
		scores = score_dict[accession][1]
		best_Evalue = min(Evalues)
		best_score = max(scores)
		n_best = scores.count(best_score)
		if n_best > 1:
			print "{0} choices for {1}".format(n_best,accession)
			##Diverge_set
			prev_index = -1
			diverge_dict = {}
			for count in range(n_best):
				prev_index = scores.index(best_score, prev_index+1)
				diverge_dict.update( { tax_node_keys[prev_index] : best_score } )
			results = diverge( path, accession, diverge_dict, tax_node )
			if results == None:
				print "Can't find a clear match for {0}".format(accession)
				ambiguous_count += 1
				#Name from fasta file', 'Genus guessed from SSUMMO','Seq. Length','Time','Score','Evalue'
				results_file.write('{0}\t{1}\n'.format(accession,path))
			else:
				accession, tax_node_key = results
				accessions_best.update( { accession : tax_node_key } )
		elif n_best == 1:
			accessions_best.update( { accession : tax_node_keys[scores.index(best_score)] } )
		elif n_best == 0:
			print "No matches for accession {0}".format(accession)
		else:
			print "n_best = {0}".format(n_best)
	return accessions_best



def scan_and_score(sequence_file_name,path,tax_node):
	"""Given the name of a sequence file, this will run hmmscan
	and return a dictionary of accessions and hmmscan scores
	as values.
	Ambiguous matches are added to key {'.' : <accessions> }
	"""
	global hmmerdir
	nodes = sorted(tax_node.keys())
	key = path[path.rfind('/')+1:]
	subprocess.call( [os.path.join(hmmerdir,'hmmscan') , '--tblout','temp_scan_table.txt','--noali', os.path.join(path,key) , sequence_file_name ], shell=False)
	#os.system('{0}/hmmscan --tblout temp_scan_table.txt --noali "{1}/{2}" "{3}"'.format( hmmerdir, path, key, sequence_file_name) )
	with file('temp_scan_table.txt','rU') as scan_output:
		temp_dict = parse_pressed_hmm(scan_output,nodes)
###         temp_dict.update( { key : [float(Evalue), float(score)] } )
	accessions_best = score(temp_dict, tax_node, path)
	return accessions_best

def SSUMMO(tax_node, path,results_node,seqs_file_name='generated'):
	global hmmerdir
	global blastversion
	global blastdir
	global blastdbdir
	global top
	informed = False
	top_len = len(top)
	tax_node_keys = sorted(tax_node.keys())
	num_tax_nodes = len(tax_node_keys)
	if tax_node_keys == []:
		print "{0} sequences assigned to {1}".format(len(results_node['accessions']),path)
		return results_node
	if path == top and seqs_file_name != 'generated':
		pass
	else:
		with file('accessions.txt','w') as accession_handle:
			accession_handle.write('\n'.join(results_node['accessions']))
		### RETRIEVE SEQUENCES FROM BLASTDB --> temp_seqs.fas
		if blastversion < 2.6:
			os.system('{0}/fastacmd -d "{1}/tempDB" -i accessions.txt -o temp_seqs.fa'.format(blastdir, blastdbdir) )
		elif blastversion == 2.6:
			os.system('{0}/blastdbcmd -db "{1}/tempDB" -entry_batch accessions.txt -out temp_seqs.fa'.format(blastdir,blastdbdir) )
		else:
			if informed == False:
				informed = True
				print "Haven't tested on a blast version this recent. Attempting to use v2.6 config settings"
			os.system('{0}/blastdbcmd -db "{1}/tempDB" -entry_batch accessions.txt -out temp_seqs.fa'.format(blastdir,blastdbdir) )
		seqs_file_name = 'temp_seqs.fa'
	with file(seqs_file_name,'rU') as seq_file:
		sequence_count = 0
		for line in seq_file.xreadlines():
			if line.startswith('>'):
				sequence_count += 1
			else:
				continue
    ##### HMMSEARCH & RETRIEVE SCORES #####
	print "hmmsearching {0} sequences against {1} HMMs in {2}.\nContents: {3}".format(sequence_count, len(tax_node_keys),path, tax_node_keys)
	accessions_best = scan_and_score(seqs_file_name,path,tax_node)
    ##### STORE RESULTS / UPDATE RESULT DICT #####
	if accessions_best != {}:
		choices = set()
		for accession, choice in accessions_best.items():
			if '|' in accession:
				accession = accession[accession.find('|')+1:]
			if choice not in results_node.keys():
				results_node.update( { choice : {'accessions' :[accession] } } )
			else:
				results_node[choice]['accessions'].append(accession)
			choices.add(choice)
    ##############################################

		for choice in choices:
			#print "number of seqs assigned to {0} is {1}".format(repr(choice),len(results_node[choice]['accessions']))
			## If node has no results, delete it.
			if len(results_node[choice]['accessions']) == 0:
				del(tax_node[choice])
			## otherwise recursively do SSuMMO on winning node.
			else:
				results_node[choice] = SSUMMO(tax_node[choice], path + '/' + choice , results_node[choice])
#				if len(tax_node[choice].keys()) != 0:
#					del(results_node[choice]['accessions'])
	else:
		print "{0} accessions assigned to node {1}".format(len(results_node['accessions']),path)
		#print tax_node.keys()
		pass
		#raise ValueError
		#num_nodes = len(nodes)
	#print "Ended up at!!! ", path[top_len:], "Score: ",best_score, "Evalue: ",best_Evalue
	return results_node
#	return path[path.rfind('/',0,path.rfind('/')-1):], best_score, best_Evalue

def prepare(in_sequence_file_name):
	"""Provide a sequence file as first command line arg and it shall be parsed one sequence
	at a time, and then SSUMMO shall be run on it...
	"""
	global results_file
	global top
	global blastversion
	global blastdir
	global blastdbdir
	seq_file_name = in_sequence_file_name
	results_file_name = seq_file_name[:seq_file_name.rfind('.')]
#	top = '/biol/people/albl500/454RDPstuff/arbDB'
	results_file = file('{0}_results.txt'.format(results_file_name),'w')
	results_file.write('%s\t%s\t%s\t%s\t%s\t%s\n' % ('Name from fasta file', 'Genus guessed from SSUMMO','Seq. Length','Time','Score','Evalue'))
	if blastversion < 2.6:
		os.system('{0}/formatdb -i "{1}" -o T -n "{2}/tempDB"'.format(blastdir, seq_file_name, blastdbdir))
	elif blastversion == 2.6:
		os.system('{0}/makeblastdb -in "{1}" -parse_seqids -hash_index -out "{2}/tempDB"'.format(blastdir, seq_file_name, blastdbdir) )
	else:
		print "Haven't tested on a blast version this recent. Attempting to use v2.6 config settings"
		os.system('{0}/makeblastdb -in "{1}" -parse_seqids -hash_index -out "{2}/tempDB"'.format(blastdir, seq_file_name, blastdbdir) )
	t = time.time()
	results = {}
	print "Formatted blastDB in {0} seconds".format(t - t0)
	try:
		results_dict = SSUMMO( load_index(), top, results,seqs_file_name = in_sequence_file_name)
	except Exception:
		results_file.close()
		raise
	t = time.time() - t0
	#results_file.write('%s\t%s\t%s\t%s\t%s\t%s\n' % (org,SSUMMO_org,len(seq),t,score, Evalue))
	#################################
#	results_file.close()
	print "processed {0} sequences in {1} seconds".format(countseqs(os.getcwd(),file_name = in_sequence_file_name),t)    

	# node = results_dict
	# node_name = 'top'
	# path = top
	# results_file.write('{0}\t{1}\t{2}\n').format(accession, node, len(seq)) ##, score, Evalue)
	# for sub_node in node.keys():
		# if sub_node == 'accessions':
			# for accession in node['accessions']:
				# results_file.write('{0}\t{1}\t{2}\t{3}\t{4}\t{5}\n'.format(accession, node_name, get_seq_len() )
		# else:
	return results_dict


if __name__ == '__main__':
	t0 = time.time()
	try:
		results_file_name = sys.argv[1][:sys.argv[1].rfind('.')]
		results_dict = prepare(sys.argv[1])
	except IndexError:
		print "Wrong number of arguments"
		raise
	print "{0} ambiguous results".format(ambiguous_count)
	### Save results files
	print "Saving results"	
	with file('{0}.pkl'.format(results_file_name),'wb') as pickle_file:
		print "Saving pickled results file to '{0}.pkl'".format(results_file_name)
		pickle.dump(results_dict,pickle_file)
	with file('{0}.xml'.format(results_file_name),'w') as write_handle:
		print "Writing phyloxml file to '{0}.xml'".format(results_file_name)
		write_xml(top,write_handle,results_dict)
	### Create bar graph index file for ITOL ###
	### CGI over to ITOL ###
	t = time.time()
	mins = int(t-t0) / 60
	print "Successfully finished in {0} mins{1} secs".format( mins, (t - t0) - (mins*60))
